from matplotlib import pyplot as plt
import networkx as nx
from networkx.drawing.nx_agraph import graphviz_layout


def extract_first_layer(processing_dag_nx: nx.DiGraph):
    ''' 抽出第一层，传入正在处理的DAG图 processing_dag_nx '''
    window_size = 0  # 定义窗口大小，只取第一层
    dag_nx: nx.DiGraph = processing_dag_nx
    front_nodes = [node for node in dag_nx.nodes()
                   if dag_nx.in_degree(node) == 0]
    cnots_in_window = []
    dag_node_id_in_window = []
    visited = set()
    has_remaining_cnot = True
    queue = [(node, 0) for node in front_nodes]  # (节点, 当前层数)
    while queue:
        current_node, depth = queue.pop(0)
        if current_node in visited or depth > window_size:
            continue
        visited.add(current_node)
        if "label" in dag_nx.nodes[current_node]:
            cnot = dag_nx.nodes[current_node]["label"]
            if len(cnot) == 2 and isinstance(cnot[0], int) and isinstance(cnot[1], int):
                cnots_in_window.append(cnot)
                dag_node_id_in_window.append(current_node)
        for successor in dag_nx.successors(current_node):
            queue.append((successor, depth + 1))
    subgraph_nodes = list(visited)
    tmp_dag = dag_nx.subgraph(subgraph_nodes).copy()
    tmp_dag_with_attrs = nx.DiGraph()
    for node in tmp_dag.nodes():
        attrs = dag_nx.nodes[node]
        tmp_dag_with_attrs.add_node(node, **attrs)  # 复制节点属性
    for u, v in tmp_dag.edges():
        attrs = dag_nx.edges[u, v]
        tmp_dag_with_attrs.add_edge(u, v, **attrs)  # 复制边属性
    if (len(cnots_in_window) == 0):
        has_remaining_cnot = False
    # draw_a_dag(tmp_dag_with_attrs) # 临时图
    return cnots_in_window, tmp_dag_with_attrs, dag_node_id_in_window, has_remaining_cnot


def remove_root_node_from_dag(graph: nx.DiGraph) -> nx.DiGraph:
    '''从DAG图中移除根节点'''
    roots = [node for node in graph.nodes() if graph.in_degree(node) == 0]
    for root in roots:
        graph.remove_node(root)
    return graph


def dag_to_json(dag: nx.DiGraph):
    print(nx.json_graph.node_link_data(dag))


def draw_a_dag(dag: nx.DiGraph):
    G = dag

    pos = graphviz_layout(G, prog='dot', args='-Grankdir=LR')
    labels = nx.get_node_attributes(G, 'label')
    edge_labels = nx.get_edge_attributes(G, 'weight')
    # plt.figure(figsize=(40, 32))
    plt.figure(figsize=(20, 16))
    nx.draw(G, pos, with_labels=True, labels=labels, node_size=2000,
            node_color='skyblue', font_size=10, font_weight='bold')
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels)
    plt.title("DAG Visualization")
    plt.show()
